import datetime
import os

import numpy as np
import pandas as pd


def load_data(X_file="./npy/data_x.npy", Y_file="./npy/data_y.npy", new=False):
    if os.path.exists(X_file) and os.path.exists(Y_file) and not new:
        X = np.load(X_file)
        Y = np.load(Y_file)
        return np.array(X), np.array(Y) ##如果特征文件已经存在，则读入特征文件
    else:
        path = './hy_round1_train_20200102'
        train_file = os.listdir(path)
        X = []
        Y = []
        for i, each in enumerate(train_file):
            if not i % 1000:
                print(i)
            each_path = os.path.join(path, each)
            x, y = read_feat(each_path, True)##调用特征向量生成函数，为每个样本生成向量
            if x is not None:
                X.append(x)
                Y.append(y)
                # print(x.shape, y.shape)
        X = np.array(X)
        Y = np.array(Y)
        print('X.shape:', X.shape, ', Y.shape:', Y.shape)
        np.save(X_file, X) ##将特征向量存入特征文件
        np.save(Y_file, Y)
        return X, Y


def read_feat(path, test_mode=False):
    MIN_LEN = 300
    LEN = 420
    N = 64

    df = pd.read_csv(path)
    df = df.iloc[::-1]

    if test_mode:   ##测试模式需要标签数据
        df['type'] = df['type'].map({'拖网': 0, '围网': 1, '刺网': 2})
        Y = np.array(df['type'].iloc[0])
    else:
        Y = None

    df['time'] = df['time'].apply(lambda x: datetime.datetime.strptime(x, "%m%d %H:%M:%S"))
    X = df[["x", "y", "速度", '方向']].apply(lambda x: (x - np.mean(x)) / np.std(x))##通过均值方差来归一化数据，利用横纵坐标、速度、方向特征来构建二维特征向量。
    for column in list(X.columns[X.isnull().sum() > 0]):
        mean_val = X[column].mean()
        X[column].fillna(mean_val, inplace=True) #利用均值进行插值，清洗空值
    X = X.dropna(axis=0)
    X = np.array(X)
    cols = X.shape[1]
    rows = X.shape[0]
    # print("cols:", cols, ", rows:", rows)
    if rows < MIN_LEN:
        return None, None
    for i in range(rows, LEN):
        b = np.zeros((1, cols))
        for j in range(N):
            b += X[i - j - 1]
        X = np.row_stack((X, b / N))
    return X[:LEN], Y


if __name__ == "__main__":
    X, Y = load_data(new=False)
    print('X.shape =', X.shape)
    print("Y.shape =", Y.shape)
    print("X[666] =\n" + str(X[666]))
    print("Y[666] =", Y[666])

    print(X[10], Y[10])
